package db.dictionaryTrie;

import beans.Word;
import dao.WordDAO;
import org.junit.Test;

import java.util.*;

/**
 * 在内存中维护一棵字典树，用于搜索提示（前缀字符串匹配）
 */
public class DictionaryTrie {

    /**
     * 全局唯一字典单词查找树
     * 也是字典树的根结点，不存储字符
     */
    private static TrieNode dictionaryTrie;

    public static void initTrie() {
        System.out.println("字典树初始化中...");
        dictionaryTrie = new TrieNode();
        List<Word> wordList = WordDAO.ME.listAllWords();
        for (Word word : wordList) {
            insert(word.getWord_string());
        }
        System.out.println("字典树初始化完成");
    }

    public static void insert(String wordStr) {
        if (wordStr == null || wordStr.isEmpty()) {
            return;
        }
        TrieNode pNode = dictionaryTrie;
        for (int i = 0;i < wordStr.length();i++) {
            TrieNode tempNode = pNode.getChildTrieNodeMap().get(wordStr.charAt(i));
            if (tempNode == null) {
                tempNode = new TrieNode();
                tempNode.setData(wordStr.charAt(i));
                pNode.getChildTrieNodeMap().put(wordStr.charAt(i), tempNode);
            }
            pNode = tempNode;
        }
        pNode.setWordEnd(true);
    }

    public static void delete(String wordStr) {
        if (wordStr == null || wordStr.isEmpty()) {
            return;
        }
        TrieNode currentNode = dictionaryTrie;
        recursiveDelete(currentNode, wordStr, 0);
    }

    /**
     * 从字典树中递归删除一个单词，如果一个节点isWordEnd值为false并且没有任何子节点，则删除该节点
     * @return 本层递归是否进行了删除节点的操作（将isWordEnd置为false也表示删除）
     * 如果下层递归没有删除节点，则本层肯定也不需要进行删除
     */
    private static boolean recursiveDelete(TrieNode currentNode, String wordStr, int i) {
        if (currentNode == null) {
            return false;
        }
        if (i == wordStr.length()) {
            currentNode.setWordEnd(false);
            return true;
        }
        TrieNode parentNode = currentNode;
        currentNode = currentNode.getChildTrieNodeMap().get(wordStr.charAt(i));
        boolean result = recursiveDelete(currentNode, wordStr, ++i);
        if (result && !currentNode.isWordEnd() && currentNode.getChildTrieNodeMap().size() == 0) {
            parentNode.getChildTrieNodeMap().remove(currentNode.getData());
            return true;
        }
        return false;
    }

    /**
     * 从字典树中查找num个以prefix为前缀的单词
     */
    public static List<String> prefixSearch(String prefix, int num) {
        List<String> wordList = new ArrayList<>();
        TrieNode pNode = findStrLastNode(prefix);
        if (pNode != null) {
            backTrackingSearch(new StringBuilder(prefix), pNode, wordList, num);
        }
        return wordList;
    }

    /**
     * 查找一个字符串的末尾在字典树中的节点位置
     */
    public static TrieNode findStrLastNode(String str) {
        if (str== null || str.isEmpty()) {
            return null;
        }
        TrieNode pNode = dictionaryTrie;
        for (int i = 0;i < str.length();i++) {
            pNode = pNode.getChildTrieNodeMap().get(str.charAt(i));
            if (pNode == null) {
                return null;
            }
        }
        return pNode;
    }

    /**
     * 利用回溯算法从某个节点开始遍历查找该节点下的所有单词
     */
    private static void backTrackingSearch(StringBuilder currentStr, TrieNode currentNode, List<String> resultList, int num) {
        if (currentNode.isWordEnd()) {
            resultList.add(currentStr.toString());
        }
        if (resultList.size() >= num) {
            return;
        }
        for (TrieNode trieNode : currentNode.getChildTrieNodeMap().values()) {
            currentStr.append(trieNode.getData());
            backTrackingSearch(currentStr, trieNode, resultList, num);
            currentStr.deleteCharAt(currentStr.length() - 1);
        }
    }

    @Test
    public static void main(String[] args) {
        initTrie();
        List<String> resultList = prefixSearch("co", 10);
        delete("co");
        resultList = prefixSearch("co", 10);
        delete("computer");
        resultList = prefixSearch("co", 10);
        System.out.println();
    }
}
